import os
from os.path import join
import numpy as np
import rospy
import rosbag
from sensor_msgs.msg import Image, CompressedImage
import cv2
import imageio
from cv_bridge import CvBridge, CvBridgeError
import re 


def uint16_to_3uint8(img):
    img1 = img.copy()
    img = np.zeros((img1.shape[0], img1.shape[1], 3)).astype(np.uint8)
    img[..., 0] = img1 & 255
    img[..., 1] = (img1 >> 8) & 255
    return img


def save_timestamp(save_path, timestamp):
    t_str = [str(t) for t in timestamp]
    with open(save_path, "w") as f:
        f.write(",".join(t_str))


def parse_rosbag(rosbag_path, save_dir, save_video=False):
    """
    RGB: uint8
    depth: uint16, unit: mm, [Note: Cannot compress 16UC1 depth image during rosbag recording!!!]
    """

    os.makedirs(save_dir, exist_ok=True)

    bag = rosbag.Bag(rosbag_path, "r")
    print(bag)

    info = bag.get_type_and_topic_info()
    print(info)

    cv_bridge = CvBridge()

    timestamps = {}
    if save_video:
        video_writers = {}
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        depth_data = {}
    VTS_data = {}

    bagMessage = bag.read_messages()
    for topic, msg, t in bagMessage:
        if topic.find("ob_camera") > -1:  # Orbbec camera: saved as sensor_msgs.Image
            print(topic, t)
            if topic.find("/color/") > -1:
                cv_image = cv_bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
            elif topic.find("/depth/") > -1:
                cv_image = cv_bridge.imgmsg_to_cv2(msg, desired_encoding="16UC1")
            else:
                raise NotImplementedError
            print(cv_image.shape, cv_image.dtype, cv_image.min(), cv_image.max())
            if not save_video:
                cv2.imwrite(join(save_dir, "vis_{}_{}.png".format(topic.replace("/", "-"), str(t))), cv_image)
            else:
                if not topic in video_writers:
                    video_writers[topic] = cv2.VideoWriter(join(save_dir, topic.replace("/", "_") + ".mp4"), fourcc, 5, (cv_image.shape[1], cv_image.shape[0]))
                    timestamps[topic] = []
                if cv_image.dtype == np.uint16:
                    cv_image = uint16_to_3uint8(cv_image)
                video_writers[topic].write(cv_image)
                timestamps[topic].append(t)
        if topic.find("d435i_camera") > -1:  # Realsense camera: saved as sensor_msgs.CompressedImage
            print(topic, t)
            if topic.find("/color/") > -1:
                cv_image = cv_bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
            elif topic.find("/aligned_depth_to_color/") > -1:
                cv_image = cv_bridge.imgmsg_to_cv2(msg, desired_encoding="16UC1")
            else:
                raise NotImplementedError
            print(cv_image.shape, cv_image.dtype, cv_image.min(), cv_image.max())
            if not save_video:
                cv2.imwrite(join(save_dir, "vis_{}_{}.png".format(topic.replace("/", "-"), str(t))), cv_image)
            else:
                if not topic in timestamps:
                    timestamps[topic] = []
                timestamps[topic].append(t)
                
                if topic.find("/color/") > -1:
                    if not topic in video_writers:
                        video_writers[topic] = cv2.VideoWriter(join(save_dir, topic.replace("/", "_") + ".mp4"), fourcc, 30, (cv_image.shape[1], cv_image.shape[0]))
                    video_writers[topic].write(cv_image)
                elif topic.find("/aligned_depth_to_color/") > -1:
                    if not topic in depth_data:
                        depth_data[topic] = []
                    depth_data[topic].append(cv_image)
                else:
                    raise NotImplementedError
        if topic.find("d455_camera") > -1:  # Realsense camera: saved as sensor_msgs.CompressedImage
            print(topic, t)
            if topic.find("/color/") > -1:
                cv_image = cv_bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
            elif topic.find("/aligned_depth_to_color/") > -1:
                cv_image = cv_bridge.imgmsg_to_cv2(msg, desired_encoding="16UC1")
            else:
                raise NotImplementedError
            print(cv_image.shape, cv_image.dtype, cv_image.min(), cv_image.max())
            if not save_video:
                cv2.imwrite(join(save_dir, "vis_{}_{}.png".format(topic.replace("/", "-"), str(t))), cv_image)
            else:
                if not topic in timestamps:
                    timestamps[topic] = []
                timestamps[topic].append(t)
                
                if topic.find("/color/") > -1:
                    if not topic in video_writers:
                        video_writers[topic] = cv2.VideoWriter(join(save_dir, topic.replace("/", "_") + ".mp4"), fourcc, 30, (cv_image.shape[1], cv_image.shape[0]))
                    video_writers[topic].write(cv_image)
                elif topic.find("/aligned_depth_to_color/") > -1:
                    if not topic in depth_data:
                        depth_data[topic] = []
                    depth_data[topic].append(cv_image)
                else:
                    raise NotImplementedError
        if topic == "/VTS_rigid":  # VTS rigid
            t = int(msg.header.stamp.secs) * 1000000000 + int(msg.header.stamp.nsecs)
            if not topic in VTS_data:
                VTS_data[topic] = []
                VTS_data["VTS_rigid_timestamp"] = []
            VTS_rigid_poses = []
            for pose in msg.poses:
                VTS_rigid_poses.append({'position': np.asarray((pose.position.x, pose.position.y, pose.position.z)),  # unit: m
                                'orientation': np.asarray((pose.orientation.w, pose.orientation.x, pose.orientation.y, pose.orientation.z))})
            VTS_data[topic].append(VTS_rigid_poses)
            VTS_data["VTS_rigid_timestamp"].append(t)
        if topic == "/rigid":  # VTS device
            t = int(msg.header.stamp.secs) * 1000000000 + int(msg.header.stamp.nsecs)
            if not topic in VTS_data:
                VTS_data[topic] = []
                VTS_data["rigid_timestamp"] = []
            rigid_poses = []
            for pose in msg.poses:
                rigid_poses.append({'position': np.asarray((pose.position.x / 100, pose.position.y / 100, pose.position.z / 100)),  # unit: m
                                'orientation': np.asarray((pose.orientation.w, pose.orientation.x, pose.orientation.y, pose.orientation.z))})
            VTS_data[topic].append(rigid_poses)
            VTS_data["rigid_timestamp"].append(t)
        if topic == "/labels":  # VTS device message
            if not topic in VTS_data:
                VTS_data[topic] = []
            VTS_data[topic].append(msg.data.split("|")[1:])
        if topic == "/pwr":  # VTS PWR tracker poses
            t = int(msg.header.stamp.secs) * 1000000000 + int(msg.header.stamp.nsecs)
            if not topic in VTS_data:
                VTS_data[topic] = []
                VTS_data["pwr_timestamp"] = []
            rigid_poses = []
            for pose in msg.poses:
                rigid_poses.append({'position': np.asarray((pose.position.x, pose.position.y, pose.position.z)),  # unit: m
                                'orientation': np.asarray((pose.orientation.w, pose.orientation.x, pose.orientation.y, pose.orientation.z))})
            VTS_data[topic].append(rigid_poses)
            VTS_data["pwr_timestamp"].append(t)
        if topic == "/pwr_ids":  # VTS PWR tracker indices
            if not topic in VTS_data:
                VTS_data[topic] = []
            VTS_data[topic].append(msg.data.split("|")[1:])
        if topic == "/joints":  # person1
            # print("ok")
            # exit(0)
            t = int(msg.header.stamp.secs) * 1000000000 + int(msg.header.stamp.nsecs)
            if not topic in VTS_data:
                VTS_data[topic] = []
                VTS_data["person1_timestamp"] = []
            joint_rawdata = []
            for pose in msg.poses:
                joint_rawdata.append({'position': np.float32([pose.position.x, pose.position.y, pose.position.z]),  # unit: cm
                                  'orientation': np.float32([pose.orientation.w, pose.orientation.x, pose.orientation.y, pose.orientation.z])})
            VTS_data[topic].append(joint_rawdata)
            VTS_data["person1_timestamp"].append(t)
        if topic == "/joints2":  # person2
            t = int(msg.header.stamp.secs) * 1000000000 + int(msg.header.stamp.nsecs)
            if not topic in VTS_data:
                VTS_data[topic] = []
                VTS_data["person2_timestamp"] = []
            joint_rawdata = []
            for pose in msg.poses:
                joint_rawdata.append({'position': np.float32([pose.position.x, pose.position.y, pose.position.z]),  # unit: cm
                                  'orientation': np.float32([pose.orientation.w, pose.orientation.x, pose.orientation.y, pose.orientation.z])})
            VTS_data[topic].append(joint_rawdata)
            VTS_data["person2_timestamp"].append(t)
    
    bag.close()

    # # check timestamps
    # for i in range(len(VTS_data["rigid_timestamp"])-1):
    #     print(i, VTS_data["rigid_timestamp"][i+1], (VTS_data["rigid_timestamp"][i+1]-VTS_data["rigid_timestamp"][i])*1e-6)
    
    if save_video:
        for topic in video_writers:
            video_writers[topic].release()
        for topic in depth_data:
            # np.save(join(save_dir, topic.replace("/", "_") + ".npy"), np.uint16(depth_data[topic]))
            os.makedirs(join(save_dir, "tmp_data"), exist_ok=True)
            for i, img in enumerate(depth_data[topic]):
                cv2.imwrite(join(save_dir, "tmp_data", str(i) + ".png"), img)
            os.system("ffmpeg -framerate 15 -f image2 -i {}/%d.png -vcodec ffv1 -level 1 -coder 1 -context 1 -g 1 -pix_fmt gray16be -qscale:v 0 {} -y -loglevel quiet".format(join(save_dir, "tmp_data"), join(save_dir, topic.replace("/", "_") + ".avi")))  # depth imgs -> avi
            os.system("rm -r {}".format(join(save_dir, "tmp_data")))

    # check data
    if "/rigid" in VTS_data:
        print("/rigid & /labels sequence lengths =", len(VTS_data["/rigid"]), len(VTS_data["/labels"]))
        N = min(len(VTS_data["/rigid"]), len(VTS_data["/labels"]))
        VTS_data["/rigid"] = VTS_data["/rigid"][:N]
        VTS_data["rigid_timestamp"] = VTS_data["rigid_timestamp"][:N]
        VTS_data["/labels"] = VTS_data["/labels"][:N]
    if "/pwr" in VTS_data:
        print("/pwr & /pwr_ids sequence lengths =", len(VTS_data["/pwr"]), len(VTS_data["/pwr_ids"]))
        N = min(len(VTS_data["/pwr"]), len(VTS_data["/pwr_ids"]))
        VTS_data["/pwr"] = VTS_data["/pwr"][:N]
        VTS_data["pwr_timestamp"] = VTS_data["pwr_timestamp"][:N]
        VTS_data["/pwr_ids"] = VTS_data["/pwr_ids"][:N]
    
    if VTS_data:
        np.savez(join(save_dir, "VTS_data.npz"), data=VTS_data)
    for topic in timestamps:
        save_timestamp(join(save_dir, topic.replace("/", "_") + "_timestamp.txt"), timestamps[topic])


if __name__ == "__main__":
    data_dir = "./rosbag_data"
    save_root = "./data"
    # data_dir = "/media/liuyun/data/HHO/multi-camera/rosbag_data"
    # save_root = "/media/liuyun/data/HHO/multi-camera/data"

    # ###################################################################
    # parse the selected file
    # ###################################################################
    # rosbag_path = join(data_dir, "2023-08-05-16-16-26.bag")
    # save_dir = join(save_root, "20230805", "test004")
    # parse_rosbag(rosbag_path, save_dir, save_video=True)
    # ###################################################################
    
    # ###################################################################
    # parse the newest file as test001, test002 etc
    # ###################################################################
    date = "2023-08-07"
    fn_latest = ""
    for fn in os.listdir(data_dir):
        if not ".bag" in fn:
            continue
        if not date in fn:
            continue
        if(fn > fn_latest):
            fn_latest = fn
    rosbag_path = join(data_dir, fn_latest)
    save_dir_home = join(save_root, date.replace("-", ""))
    fn_test_max = ""
    if not os.path.exists(save_dir_home):
        os.makedirs(save_dir_home)
    for fn in os.listdir(save_dir_home):
        if not "test" in fn:
            continue
        if(fn > fn_test_max):
            fn_test_max = fn
    match = re.search(r'\d{3,}$',fn_test_max)
    num = 0
    if  not match is None:
        num = int(match.group()) + 1
    save_dir = join(save_dir_home, "test" + str(num).zfill(3))
    print(save_dir)
    print(rosbag_path)
    parse_rosbag(rosbag_path, save_dir, save_video=True)
    # ###################################################################
    

    # ###################################################################
    # parse all file 
    # ###################################################################
    # date = "2023-08-02"
    # fns = []
    # for fn in os.listdir(data_dir):
    #     if not ".bag" in fn:
    #         continue
    #     if not date in fn:
    #         continue
    #     fns.append(fn)
    # fns.sort()
    # for (i, fn) in enumerate(fns):
    #     rosbag_path = join(data_dir, fn)
    #     save_dir = join(save_root, date.replace("-", ""), str(i).zfill(3))
    #     parse_rosbag(rosbag_path, save_dir, save_video=True)
    # ###################################################################
